from typing import Sized
from matplotlib.pyplot import close, plot
import pandas as pd
import backtrader as bt
from backtrader_plotting import Bokeh
from backtrader_plotting.schemes import Tradimo
from my_plot import *
from py_mysql import *
import numpy as np

class PandasData_more(bt.feeds.PandasData):
    pass

# 创建策略继承bt.strategy
class TestStrategy(bt.Strategy):
    params = (
        ('maperiod', 15),
        # 判断是否输出该日志
        ('printlog', False),
    )
    # sma_periodArr = [30,50,100]

    # 初始化部分指标

    def __init__(self):
        self.dataclose = self.datas[0].close
        # 买入价格和手续费
        self.buyprice = None
        self.buycomm = None
        # 加入均线指标
        self.MA30 = bt.indicators.MovingAverageSimple(self.datas[0].close, period=30)
        self.MA50 = bt.indicators.MovingAverageSimple(self.datas[0].close, period=50)
        self.MA100 = bt.indicators.MovingAverageSimple(self.datas[0].close, period=100)

        # 下单数量
        self.volSize = 1
        # 持仓5分钟后清仓
        self.timeSum = 0
        # 状态值
        self.status = None
        # 当前状况
        self.buyStatus = False
        self.sellStatus = False
        # 止盈止损
        self.ZY_price = 0
        self.ZS_price = 0
        # 止盈止损比例
        self.ZY_size = 1
        self.ZS_size = 1


    # 策略核心代码
    def next(self):
        # 当前日期
        self.time = str(bt.num2date(self.lines.datetime[0]))
        self.positionSize = self.position.size

        # 当前空仓
        if not self.position:
            if self.MA30[0] > self.MA50[0] and self.MA50[0] > self.MA100[0] and (not(self.MA30[-1] > self.MA50[-1] and self.MA50[-1] > self.MA100[-1])):
                if self.dataclose[0] < self.MA30[0] and self.dataclose[-1] > self.MA30[-1]:
                    self.buyStatus = True
                if self.dataclose[0] < self.MA50[0]:
                    self.buyStatus = False

                if self.buyStatus and self.dataclose[0] > self.MA30[0]:
                    self.buy(size=self.volSize)
                    self.status = 'buy'
                    self.buyStatus = False
                    sizeNUm = self.dataclose[0] - self.MA100[0]
                    self.ZY_price = self.dataclose[0] + self.ZY_size*sizeNUm
                    self.ZS_price = self.dataclose[0] - self.ZS_size*sizeNUm
                    print('出现买入:---，止盈:{}---止损:{}'.format(self.ZY_price,self.ZS_price))

                    
            if self.MA30[0] < self.MA50[0] and self.MA50[0] < self.MA100[0] and (not(self.MA30[-1] < self.MA50[-1] and self.MA50[-1] < self.MA100[-1])):
                if self.dataclose[0] > self.MA30[0] and self.dataclose[-1] < self.MA30[-1]:
                    self.sellStatus = True
                if self.dataclose[0] > self.MA50[0]:
                    self.sellStatus = False
                
                if self.sellStatus and self.dataclose[0] < self.MA30[0]:
                    self.sell(size=self.volSize)
                    self.status = 'sell'
                    self.sellStatus = False
                    sizeNUm = self.MA100[0] - self.dataclose[0]
                    self.ZY_price = self.dataclose[0] - self.ZY_size*sizeNUm
                    self.ZS_price = self.dataclose[0] + self.ZS_size*sizeNUm
        else:
            if self.status == 'buy':
                if self.dataclose[0] > self.ZY_price or self.dataclose[0] < self.ZS_price:
                    self.close()
                    # self.status = None
                    
            if self.status == 'sell':
                if self.dataclose[0] > self.ZS_price or self.dataclose[0] < self.ZY_price:
                    self.close()
                    # self.status = None
            

    # 日志输出
    def log(self, txt, dt=None, doprint=False):
        if self.params.printlog or doprint:
            # 记录策略的执行日志
            dt = dt or self.datas[0].datetime.date(0)
            print(f'{dt.isoformat()},{txt}')

    # 订单状态通知，买入卖出都是下单
    def notify_order(self, order):
        # 提交了/接受了,  买/卖订单什么都不做
        if order.status in [order.Submitted, order.Accepted]:
            return
        # 检查一个订单是否完成
        # 注意:当资金不足时，broker会拒绝订单
        if order.status in [order.Completed]:
            if order.isbuy():
                print('日期：{}---买入价格:{}---买入手续费{}'.format(self.time,order.executed.price,
                      order.executed.comm))
                self.buyprice = order.executed.price
                self.buycomm = order.executed.comm
            elif order.issell():
                print('日期：{}---卖出价格:{}---卖出手续费{}'.format(self.time,order.executed.price,
                      order.executed.comm))
            self.bar_executed = len(self)
        elif order.status in [order.Canceled, order.Margin, order.Rejected]:
            self.log('订单取消/保证金不足/拒绝', doprint=False)
        # 其他状态记录为：无法挂单
        self.order = None

    # 交易状态通知，一买一卖算交易（交易净利润）
    def notify_trade(self, trade):
        if not trade.isclosed:
            return

    # 策略结束时，多用于参数调优
    def stop(self):
        print(self.positionSize)
        self.log('(MA均线： %2d日) 期末总资金 %.2f' %
                 (self.params.maperiod, self.broker.getvalue()), doprint=False)


def main(startcash=1000000, com=0.0003, qts=100):

    # 创建主控制器
    cerebro = bt.Cerebro()
    # 导入策略参数寻优
    cerebro.addstrategy(TestStrategy)
    # cerebro.optstrategy(TestStrategy,maperiod=range(10, 15))
    # 获取数据
    # df = writeExcl()  
    # df.index = pd.to_datetime(df.date)
    # print(df)
    # df = df[['open', 'high', 'low', 'close', 'volume']]
    # # 将数据加载至回测系统
    # data = bt.feeds.PandasData(dataname=df)
    # cerebro.adddata(data)
    # 获取数据
    query_db = Mysql_search()
    df = query_db.get_one(['RB'],'2020-06-01','2020-12-31')
    for item in df:
        df2 = df[item]
        print(df2)
        df2.index = pd.to_datetime( df2.index, utc= True)
        data = PandasData_more(
            dataname = df2,
            timeframe=bt.TimeFrame.Minutes)
        cerebro.adddata(data)
    # broker设置资金、手续费
    cerebro.broker.setcash(startcash)
    cerebro.broker.setcommission(commission=com)
    # 设置买入设置，策略，数量
    cerebro.addsizer(bt.sizers.FixedSize, stake=qts)
    # 以发出信号当日收盘价成交
    cerebro.broker.set_coc(True)

    print('期初总资金: %.2f' % cerebro.broker.getvalue())
    cerebro.addanalyzer(bt.analyzers.TimeReturn, _name='_TimeReturn')
    result = cerebro.run()
    print('期末总资金: %.2f' % cerebro.broker.getvalue())

    # custom_plot(result)
    cerebro.plot()
    
    b = Bokeh(style='bar', tabs='multi', scheme=Tradimo())
    # cerebro.plot(b)

# def writeExcl():
#     df = pd.read_excel(r'D:\通达信\T0002\export2\IL8.xls')
#     df['code'] = 'IL8'
#     return df



def writeExcl():
    df = pd.read_excel(r'D:\Downloads\RBHour.xls')
    # df['code'] = 'MAL8'
    return df

main()
